"""
pytorch数据集 需要重写dataset， sampler之类的
"""
import os
import torch
import numpy as np
from torch import tensor
from torch.utils.data import DataLoader, Dataset  # 不是 dataloader 和 dataset
from torch.utils.data.sampler import Sampler 


class NormalDistDataSet(Dataset):
    """ class to handle dataset """

    def __init__(self, dataset_path='../Data', set='training', transform=None):
        super().__init__()

        self.transform = transform
        self.set = set
        if self.set == 'training':
            data = 'multivariate_normal_data.npy'
        elif self.set == 'validation':
            data = 'multivariate_normal_validation.npy'
        else:
            raise ValueError('Unknown set for data: ', self.set)

        data = np.load(os.path.join(dataset_path, data))
        self.data = self._handleNpyFile(data)
        self.N = len(self.data)
    
    def __len__(self): # 必须重写的方法
        return self.N
    
    def __getitem__(self, index): # 必须重写的方法
        xy, z = self.data[index]
        if self.transform is not None:
            xy = self.transform(xy)
        xy = torch.tensor(xy)
        z = torch.tensor(z)
        return xy, z
        
    def _handleNpyFile(self, data): 
        xy = data[:, 0:2]
        z = data[:, 2]
        z = np.expand_dims(z, axis=1)
        return list(zip(xy, z))

class NormalDistSampler(Sampler):
    """
    如果需要使用自己选取数据的方式，则需要重写这个类
    """
    def __init__(self, data_source: NormalDistDataSet) -> None:
        super().__init__(data_source)


if __name__ == '__main__':
    data = 'multivariate_normal_data.npy'
    data = np.load( os.path.join('../Data', data) )
    xy = data[:, 0:2]
    z = data[:, 2]
    z = np.expand_dims(z, axis=1)
    a = list(zip(xy, z))
    
    print(z.shape)